package llm_server

import (
	"context"
	"fmt"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/utils"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"io"
	"strings"
	"xiaozhi/config"
)

// ConvertMessages converts []config.Message to []*model.ChatCompletionMessage
func ConvertMessages(configMessages []config.Message) []*model.ChatCompletionMessage {
	modelMessages := make([]*model.ChatCompletionMessage, len(configMessages))
	for i, msg := range configMessages {
		modelMessages[i] = &model.ChatCompletionMessage{
			Role: msg.Role,
			Content: &model.ChatCompletionMessageContent{
				StringValue: volcengine.String(msg.Content),
			},
		}
	}
	return modelMessages
}

func (llm *LLM) OnDouBaoMessage(prompt string, systemPrompt string, history []config.Message) (string, error) {
	// 构造消息
	message := []*model.ChatCompletionMessage{
		{
			Role: model.ChatMessageRoleSystem,
			Content: &model.ChatCompletionMessageContent{
				StringValue: volcengine.String(systemPrompt),
			},
		},
		{
			Role: model.ChatMessageRoleUser,
			Content: &model.ChatCompletionMessageContent{
				StringValue: volcengine.String(prompt),
			},
		},
	}
	// 添加历史消息
	if history != nil {
		message = append(message, ConvertMessages(history)...)
	}
	client := arkruntime.NewClientWithApiKey(
		config.NewDouBaoLLMConfig().ApiKey,
	)
	ctx := context.Background()
	var contentBuilder strings.Builder
	fmt.Println("----- streaming request -----")
	req := model.ChatCompletionRequest{
		Model:    config.NewDouBaoLLMConfig().Model,
		Messages: message,
	}
	// 创建流式请求
	stream, err := client.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("stream chat error: %v\n", err)
		return "", err
	}
	// 释放资源
	defer func(stream *utils.ChatCompletionStreamReader) {
		err := stream.Close()
		if err != nil {
			_ = fmt.Errorf("stream chat close error: %v\n", err)
		}
	}(stream)
	// 读取流
	for {
		rec, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("\n----- streaming end -----")
			return contentBuilder.String(), nil
		}
		if err != nil {
			fmt.Printf("Stream chat error: %v\n", err)
			return "", err
		}

		if len(rec.Choices) > 0 {
			fmt.Print(rec.Choices[0].Delta.Content)
			contentBuilder.WriteString(rec.Choices[0].Delta.Content)
		}

	}
}
